import torch
import torch.nn as nn
import numpy as np

DEVICE = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")


class AddGaussianNoise(nn.Module):
    def __init__(self, mean=0.0, variance=0.15, amplitude=1.0, seed=479):
        super(AddGaussianNoise, self).__init__()

        self.mean = mean            # 均值
        self.variance = variance    # 方差
        self.amplitude = amplitude  # 幅度
        self.rng = np.random.RandomState(seed)

    def forward(self, x):
        noise = self.amplitude * self.rng.normal(loc=self.mean, scale=self.variance, size=x.shape)
        x += torch.tensor(noise, dtype=torch.float, device=DEVICE)
        return x


class FGM(nn.Module):
    def __init__(self, model):
        super(FGM, self).__init__()
        self.model = model
        self.backup = {}

    def attack(self, epsilon=1., emb_name='emb'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='emb'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}


class PGD(nn.Module):
    def __init__(self, model):
        super(PGD, self).__init__()
        self.model = model
        self.backup = {}
        self.grad_backup = {}
        self.attack_num = 10  # 10 3

    def attack(self, epsilon=1., alpha=0.1, emb_name='emb', is_first_attack=False):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                if is_first_attack:
                    self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data, epsilon)

    def project(self, param_name, param_data, epsilon):
        r = param_data - self.backup[param_name]
        if torch.norm(r) > epsilon:
            r = epsilon * r / torch.norm(r)
        return self.backup[param_name] + r

    def restore(self, emb_name='emb'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

    def backup_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.grad_backup[name] = param.grad.clone()

    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.grad = self.grad_backup[name]
